import pydash
from funasr import AutoModel
from funasr.utils.postprocess_utils import rich_transcription_postprocess


def _get_origin_model():
    # TODO: 解决模型下载的问题
    model_dir = "iic/SenseVoiceSmall"
    return AutoModel(
        model=model_dir,
        vad_model="fsmn-vad",
        vad_kwargs={"max_single_segment_time": 30000},
        device="cuda:0",
    )


_get_model = pydash.memoize(_get_origin_model)


def do_funasr(input):
    # en
    res = _get_model().generate(
        input=input,
        cache={},
        language="auto",  # "zn", "en", "yue", "ja", "ko", "nospeech"
        use_itn=True,
        batch_size_s=60,
        merge_vad=True,  #
        merge_length_s=15,
    )
    text = rich_transcription_postprocess(res[0]["text"])
    return text
